const burnin = 2499

function maketraceplots(mylabels,myfuns,burnin=burnin,history=history)
    myplots = Any[]
    hist = history[save_params.>burnin]
    iter = save_params[save_params.>burnin]
    for (n,myname) in enumerate(mylabels)
        plt = plot()
        for t=1:4
            betas = [myfuns[n](h,t) for h in hist]
            plot!(plt,iter,betas,label="Type $t",legend=(n==1 ? :left : false))
        end
        title!(plt,myname,titlefontsize=8,guidefontsize=6) #ylabel!(plt,myname)
        n>=length(mylabels)-1 && xlabel!(plt,"iteration")
        push!(myplots,plt)
    end
    return myplots
end

##########################
# traceplots
##########################

#match terms
interactlabels = [xj*" x "*xi for xj in ["STEM", "Humanities"], xi in ["Math", "Verbal"]]
scorelabels = ["Score*Teaching", "Score^2*Teaching","Score^3*Teaching"]
ijlabels = ["Price", "Same City", interactlabels...,scorelabels...]
# lab = [ijlabels; "Off-Platform"]
# funs = [[(h,t)->h.beta_ij[n,t] for n=1:length(ijlabels)]...,(h,t)->h.beta_fixed[1,t]]
ijplots = maketraceplots(ijlabels, [(h,t)->h.beta_ij[n,t] for n=1:length(ijlabels)])
plt=plot(ijplots..., )#layout=(4,2))
savefig(plt,figurepath*"/traceplot_ijterms.png")

#frictions
alphalabels = ["G25", "G8 On", "G8 Off", "Local"]
myfuns = [(h,t)->h.alpha[n,t] for n=1:length(alphalabels)]
alphaplots = maketraceplots(alphalabels,myfuns)
plt=plot(alphaplots...)
savefig(plt,figurepath*"/traceplot_frictions.png")

#1st oo mean params
oolabels = ["Constant", "Math", "Verbal", "Big City", "Current Cohort", "1(2011)", "1(2012)", "Scholarship Amount", "Score2","Score3"]
myfuns = [(h,t)->h.betaOO0[n,t] for n=1:length(oolabels)]
plts = maketraceplots(oolabels,myfuns)
plt = plot(plts...,layout=grid(5,2))
savefig(plt,figurepath*"/traceplot_firstOutsideOptionMean.png")

#2nd oo mean params
myfuns = [(h,t)->h.betaOO1[n,t] for n=1:length(oolabels)]
plts = maketraceplots(oolabels,myfuns)
plt = plot(plts...,layout=grid(5,2))
savefig(plt,figurepath*"/traceplot_secondOutsideOptionMean.png")

#oo sd
varlabels = ["sd of o.o. 1", "sd of o.o. 2"]
myfuns = [(h,t)->sqrt(h.sigsqOO0[t]), (h,t)->sqrt(h.sigsqOO1[t])]
plts = maketraceplots(varlabels,myfuns)
plt=plot(plts...,layout=grid(2,1))
savefig(plt,figurepath*"/traceplot_oovariance.png")

#outcomes
outcomelabels = ["utility shock"; oolabels; ijlabels]
myfuns_outcomes = [(h,t)->h.betaOutcome[n,t] for n=1:length(outcomelabels)]
plts = maketraceplots(outcomelabels,myfuns_outcomes)
plt = plot(plts[1:6]...,layout=(3,2))
savefig(plt,figurepath*"/traceplot_outcome_1.png")
plt = plot(plts[7:12]...,layout=(3,2))
savefig(plt,figurepath*"/traceplot_outcome_2.png")
plt = plot(plts[13:end]...,)
savefig(plt,figurepath*"/traceplot_outcome_3.png")

##########################
# parameter estimates
##########################
history_short = history[save_params .> burnin]
function outcomeParametersTable(f=stdout)
    println(f,"\\begin{tabular}{lcccc}")
    println(f,"\\hline\\hline")
    println(f,"Parameters & Male Private & Male Public & Female Private & Female Public \\\\ \\hline")
    println(f,"& & & & \\\\")
    #
    println(f,"Production Function & & & & \\\\")
    (cc,_,_) = dataset[2012]
    noo = size(cc.Xoo,2)
    nmatch = size(cc.Xij[1],2)
    funsU = [h->h.betaOutcome[n,:] for n=1:1]
    funsOO = [h -> h.betaOutcome[n+1,:] for n=1:noo]
    funsIJ = [h -> h.betaOutcome[n+noo+1,:] .+ h.beta_ij[n,:].*h.betaOutcome[1,:] for n=1:nmatch]
    outcomefuns = [funsU; funsOO; funsIJ]
    for (lab,fun) in zip(outcomelabels,outcomefuns )
        minirow(lab,fun,f)
    end
    println(f," & & & & \\\\ \\hline \\hline")
    println(f,"\\end{tabular}")
end

function selectedParametersTable(f=stdout)
    println(f,"\\begin{tabular}{lcccc}")
    println(f,"\\hline\\hline")
    println(f,"Parameters & Male Private & Male Public & Female Private & Female Public \\\\ \\hline")
    println(f,"& & & & \\\\")
    #
    println(f,"Aftermarket frictions (\$\\alpha\$) & & & & \\\\")
    minipanel(:alpha,alphalabels,f)
    #
    println(f," & & & & \\\\ \\hline \\hline")
    println(f,"\\end{tabular}")
end


function insideParametersTable(f=stdout)
    println(f,"\\begin{tabular}{lcccc}")
    println(f,"\\hline\\hline")
    println(f,"Parameters & Male Private & Male Public & Female Private & Female Public \\\\ \\hline")
    println(f,"& & & & \\\\")
    #
    println(f,"Preferences (\$\\psi^o\$) & & & & \\\\")
    minipanel(:beta_ij,ijlabels,f)
    #minirow("Off-Platform",mean(h.beta_fixed[1,:] for h in history_short), [std(h.beta_fixed[1,t] for h in history_short) for t=1:4],f)
    #
    println(f,"Aftermarket frictions (\$\\alpha\$) & & & & \\\\")
    minipanel(:alpha,alphalabels,f)
    #
    println(f,"SD of program FE & & & & \\\\")
    sdFE = [vec( std(h.beta_fixed[1:end,:],dims=1)) for h in history_short]
    minirow("\$\\sigma_{FE}\$",mean(sdFE),[std([s[t] for s in sdFE]) for t=1:4],f)
    #
    println(f,"RC covariance matrix (\$\\psi^u\$) & & & & \\\\")
    stem = [h.Sigma_rc[1,1,:] for h in history_short]
    minirow("STEM",mean(stem),[std([s[t] for s in stem]) for t=1:4],f)
    hum = [h.Sigma_rc[2,2,:] for h in history_short]
    minirow("Humanities",mean(hum),[std([s[t] for s in hum]) for t=1:4],f)
    mycov = [h.Sigma_rc[1,2,:] for h in history_short]
    minirow("Humanities vs STEM (\$\\rho\$)",mean(mycov),[std([s[t] for s in mycov]) for t=1:4],f)
    #
    println(f," & & & & \\\\ \\hline \\hline")
    println(f,"\\end{tabular}")
end

#[ones(size(xloc)) Sinorm(:,1:2) xloc current_cohort is2012 AidAmount(:,1)];
function outsideParametersTable(f=stdout)
    println(f,"\\begin{tabular}{lcccc}")
    println(f,"\\hline\\hline")
    println(f,"Parameters & Male Private & Male Public & Female Private & Female Public \\\\ \\hline")
    println(f,"& & & & \\\\")
    #
    println(f,"First Outside Option (\$\\gamma^0\$) & & & & \\\\")
    minipanel(:betaOO0,oolabels,f)
    minirow("\$\\sigma_{0,0}\$",mean(sqrt.(h.sigsqOO0) for h in history_short), [std(sqrt(h.sigsqOO0[t]) for h in history_short) for t=1:4],f)
    #
    println(f,"Second Outside Option (\$\\gamma^1\$) & & & & \\\\")
    minipanel(:betaOO1,oolabels,f)
    minirow("\$\\sigma_{0,1}\$",mean(sqrt.(h.sigsqOO1) for h in history_short), [std(sqrt(h.sigsqOO1[t]) for h in history_short) for t=1:4],f)
    #
    println(f," & & & & \\\\ \\hline \\hline")
    println(f,"\\end{tabular}")
end

open(tablepath*"/mcmc_estimation_parameters_insideprefs.tex","w") do f
    insideParametersTable(f)
end

open(tablepath*"/mcmc_estimation_parameters_outsideprefs.tex","w") do f
    outsideParametersTable(f)
end

open(tablepath*"/mcmc_estimation_parameters_outcomes.tex","w") do f
    outcomeParametersTable(f)
end

open(tablepath*"/mcmc_estimation_parameters_selected.tex","w") do f
    selectedParametersTable(f)
end

#program FE
let
    programFE = cat([h.beta_fixed[2:end,:] for h in history_short]...,dims=3)
    programFE_m = reshape(mean(programFE,dims=3), (size(programFE,1),size(programFE,2)))
    programFE_sd = reshape(std(programFE,dims=3), (size(programFE,1),size(programFE,2)))
    #code up G8 indicators
    cc = dataset[2012][1]
    G8 = counterfactuals[2012][1].G8
    G8_programFE = zeros(Int,size(programFE,1))
    for j=1:length(G8)
        ind = findfirst(cc.Xfixed[j,2:end].==1)
        isa(ind,Int) && (G8_programFE[ind] = G8[j])
    end
    #inds = sortperm(vec(mean(programFE_m,dims=2)))
    plts_fe_bytype = Any[]
    for t=1:4
        inds = sortperm(programFE_m[:,t])
        plt = scatter(1:length(inds),programFE_m[inds,t],yerror=programFE_sd[inds,t],legend=:false,ms=0)
        xlabel!(plt,typenames[t])
        ylabel!(plt,"Program Effects")
        push!(plts_fe_bytype,plt)
    end
    plot_fe = plot(plts_fe_bytype...,layout=grid(2,2))
    savefig(plot_fe,figurepath*"/programFE1.png")
    plts_fe_comparison = Any[]
    for (t,t2) in zip([2,3,2,3],[1,1,4,4])
        plt2 = scatter(programFE_m[:,t], programFE_m[:,t2],
            group=G8_programFE, label=["G25" "G8"], legend= ((t==2 && t2==1) ? :bottomright : false))
        xlabel!(plt2,typenames[t])
        ylabel!(plt2,typenames[t2])
        push!(plts_fe_comparison, plt2)
    end
    plot_comparisons = plot(plts_fe_comparison...,layout=grid(2,2))
    savefig(plot_comparisons,figurepath*"/programFE2.png")
end
